# ------------------------------------------------------------------------
# QAHOI
# Copyright (c) 2021 Junwen Chen. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# ------------------------------------------------------------------------

"""
QAHOI model and criterion classes.
"""
import torch
from torch import nn
import math

from util.misc import inverse_sigmoid

from .module.modules import MLP

from .backbones.backbone import build_backbone
#from .backbones.swin_backbone import build_backbone
from .necks.deformable_transformer import build_transformer
from .matcher.matcherhoi import build_matcher
from .losses.setcriterionhoi import SetCriterionHOI
from .postprocess.postprocesshoi import PostProcessHOI


from .deformable_detr import DeformableDETR


def build(args):
    device = torch.device(args.device)
    backbone = build_backbone(args)
    transformer = build_transformer(args)
    model = QAHOI(
        backbone,
        transformer,
        num_classes=args.num_obj_classes,
        num_verb_classes=args.num_verb_classes,
        num_queries=args.num_queries,
        num_feature_levels=args.num_feature_levels,
        aux_loss=args.aux_loss,
        args=args,
    )
    matcher = build_matcher(args, 'focal_loss')
    weight_dict = get_weight_dict(args)
    losses = get_losses(args)
        
    criterion = SetCriterionHOI(args, matcher=matcher, weight_dict=weight_dict, 
                                losses=losses, loss_type='focal_loss')
    criterion.to(device)
    postprocessors = {'hoi': PostProcessHOI(args, trans_type='ddetr')}

    return model, criterion, postprocessors



class QAHOI(DeformableDETR):
    """ This is the Deformable DETR module that performs object detection """
    def __init__(self, backbone, transformer, num_classes, num_verb_classes, num_queries, num_feature_levels,
                 args=None, aux_loss=True):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the deformable transformer architecture. See deformable_transformer.py
            num_classes: number of object classes
            num_queries: number of object queries, ie detection slot.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super(QAHOI,self).__init__(backbone, transformer, num_classes, num_queries, num_feature_levels, aux_loss, args, num_verb_classes)


    def build_head(self, num_classes, num_verb_classes):
        hidden_dim = self.hidden_dim
        self.class_embed = nn.Linear(hidden_dim, num_classes)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        self.verb_class_embed = nn.Linear(hidden_dim, num_verb_classes)
        self.sub_bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)

        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        self.class_embed.bias.data = torch.ones(num_classes) * bias_value
        self.verb_class_embed.bias.data = torch.ones(num_verb_classes) * bias_value     
        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
        nn.init.constant_(self.sub_bbox_embed.layers[-1].bias.data, 0)
        nn.init.constant_(self.sub_bbox_embed.layers[-1].bias.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
        nn.init.constant_(self.sub_bbox_embed.layers[-1].bias.data[2:], -2.0)
        
        num_pred = self.transformer.decoder.num_layers
        self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
        self.verb_class_embed = nn.ModuleList([self.verb_class_embed for _ in range(num_pred)])
        self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
        self.sub_bbox_embed = nn.ModuleList([self.sub_bbox_embed for _ in range(num_pred)])

        self.transformer.decoder.sub_bbox_embed = None

    def head_forward(self, outputs):
        hs = outputs['out_query']
        init_reference = outputs['init_reference']
        inter_references = outputs['inter_references']
        
        outputs_classes = []
        outputs_verb_classes = []
        outputs_coords = []
        outputs_sub_coords = []
        
        for lvl in range(hs.shape[0]):
            outputs_class = self.class_embed[lvl](hs[lvl]) 
            outputs_classes.append(outputs_class)
                      
            outputs_verb_class = self.verb_class_embed[lvl](hs[lvl])
            outputs_verb_classes.append(outputs_verb_class)            
            
            tmp = self.bbox_embed[lvl](hs[lvl])
            tmp_sub = self.sub_bbox_embed[lvl](hs[lvl])

            if lvl == 0:
                reference = init_reference
            else:
                reference = inter_references[lvl - 1]
            reference = inverse_sigmoid(reference)
            if reference.shape[-1] == 4:
                tmp += reference
                tmp_sub += reference
            else:
                assert reference.shape[-1] == 2
                tmp[..., :2] += reference
                tmp_sub[..., :2] += reference
            outputs_coord = tmp.sigmoid()
            outputs_coords.append(outputs_coord)
            
            
            outputs_sub_coord = tmp_sub.sigmoid()
            outputs_sub_coords.append(outputs_sub_coord)
            #reference = reference.sigmoid()
            
        outputs_class = torch.stack(outputs_classes)
        outputs_verb_class = torch.stack(outputs_verb_classes)
        outputs_coord = torch.stack(outputs_coords)
        outputs_sub_coord = torch.stack(outputs_sub_coords)
        
        out = {'pred_obj_logits': outputs_class[-1], 
               'pred_obj_boxes': outputs_coord[-1],
               'pred_verb_logits': outputs_verb_class[-1], 
               'pred_sub_boxes': outputs_sub_coord[-1],
               
               'pred_obj_logits_cascade': outputs_class,
               'pred_verb_logits_cascade': outputs_verb_class,
               'pred_sub_boxes_cascade': outputs_sub_coord,
               'pred_obj_boxes_cascade': outputs_coord,
               }
        if self.aux_loss and self.training:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord, outputs_verb_class, outputs_sub_coord)

        return out
      

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord, outputs_verb_class, outputs_sub_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_obj_logits': a, 
                 'pred_obj_boxes': b, 
                 'pred_verb_logits': c, 
                 'pred_sub_boxes': d}
                for a, b, c, d in zip(outputs_class[:-1], outputs_coord[:-1], outputs_verb_class[:-1], outputs_sub_coord[:-1])]


def get_weight_dict(args):
    weight_dict = {}
    weight_dict['loss_obj_ce'] = args.obj_loss_coef
    weight_dict['loss_verb_ce'] = args.verb_loss_coef
    
    weight_dict['loss_sub_bbox'] = args.bbox_loss_coef
    weight_dict['loss_obj_bbox'] = args.bbox_loss_coef
    weight_dict['loss_sub_giou'] = args.giou_loss_coef
    weight_dict['loss_obj_giou'] = args.giou_loss_coef
    # TODO this is a hack 
    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers_hopd - 1):
            #   为中间层输出的loss也加上对应权重
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)
    return weight_dict

def get_losses(args):
    losses = ['obj_labels', 'verb_labels', 'sub_obj_boxes', 'obj_cardinality']
    return losses